/******************************************************************/
/*                                                                */
/*  Winpooch : Windows Watchdog                                   */
/*  Copyright (C) 2004-2005  Benoit Blanchon                      */
/*                                                                */
/*  This program is free software; you can redistribute it        */
/*  and/or modify it under the terms of the GNU General Public    */
/*  License as published by the Free Software Foundation; either  */
/*  version 2 of the License, or (at your option) any later       */
/*  version.                                                      */
/*                                                                */
/*  This program is distributed in the hope that it will be       */
/*  useful, but WITHOUT ANY WARRANTY; without even the implied    */
/*  warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR       */
/*  PURPOSE.  See the GNU General Public License for more         */
/*  details.                                                      */
/*                                                                */
/*  You should have received a copy of the GNU General Public     */
/*  License along with this program; if not, write to the Free    */
/*  Software Foundation, Inc.,                                    */
/*  675 Mass Ave, Cambridge, MA 02139, USA.                       */
/*                                                                */
/******************************************************************/


/******************************************************************/
/* Build configuration                                            */
/******************************************************************/

#define SPYDLL_BUILD_DLL

#define TRACE_LEVEL	2 /* warning */


/******************************************************************/
/* Includes                                                       */
/******************************************************************/

// module's interface
#include "HookSys.h"

// standard headers
#include <tchar.h>
#include <shlwapi.h> // functions PathXxx

// project's headers
#include "Assert.h"
#include "SpyDll.h"
#include "SpyDllInt.h"
#include "Hooks.h"
#include "Strlcpy.h"
#include "Trace.h"


/******************************************************************/
/* Internal functions                                             */
/******************************************************************/

BOOL _HookSys_ExtractPathFromCmdLine (LPTSTR szPath, LPCTSTR szCurDir, LPCTSTR szCmdLine, 
				      BOOL bRemoveArgs, BOOL bSearchPath) ;


/******************************************************************/
/* Exported function : CreateProcessA                             */
/******************************************************************/

BOOL WINAPI Spy_CreateProcessA (LPCSTR			szAppName,
				LPSTR			szCmdLine,
				LPSECURITY_ATTRIBUTES	pProcessAttr,
				LPSECURITY_ATTRIBUTES	pThreadAttr,
				BOOL			bInheritHandles,
				DWORD			dwCreationFlags,
				PVOID			pEnvironment,
				LPCSTR			szCurrentDirectory,
				LPSTARTUPINFOW		pStartupInfo,
				LPPROCESS_INFORMATION	pProcessInformation) 
{
  PROC			pfnStub = Hooks_GetStubAddress (HOOKS_CREATEPROCESSA) ;

  WORD			nReaction ;
  BOOL			bSuccess ;
  DWORD			dwServerResponse ;
  DWORD			dwSaveErrorCode ;
  SDNCREATINGPROCESS	sdncp ;
  SDNPROCESSCREATED	sdnpc ;
  PROCESS_INFORMATION	pi ;

  BOOL			bAppNameSpecified ;
  LPTSTR		tszCmdLine, tszCurDir ;
  TCHAR			tszPath[MAX_PATH] ;

  TRACE ;

  if( szAppName==NULL && szCmdLine==NULL )
    TRACE_WARNING (TEXT("szAppName==NULL and szCmdLine==NULL\n"));

  bAppNameSpecified = szAppName != NULL ;
  
  // get command line in unicode format
  tszCmdLine = SpyInt_AsciiToUnicode (bAppNameSpecified ? szAppName : szCmdLine) ;

  // get directory in unicode format
  tszCurDir = SpyInt_AsciiToUnicode (szCurrentDirectory) ;
   
  // extract program path 
  _HookSys_ExtractPathFromCmdLine (tszPath, tszCurDir, tszCmdLine, 
				   !bAppNameSpecified, !bAppNameSpecified) ;

  // apply filter
  nReaction = SpyInt_CatchCall (REASON_SYS_EXECUTE, TEXT("s"), tszPath) ;

  // if filter reject call, just return false
  if( nReaction == RULE_REJECT ) {
    free (tszCurDir) ;
    free (tszCmdLine) ;
    SetLastError (ERROR_ACCESS_DENIED) ;
    return FALSE ; 
  }

  // if filter reject call, just return true
  if( nReaction == RULE_FEIGN ) {
    free (tszCurDir) ;
    free (tszCmdLine) ;
    SetLastError (ERROR_SUCCESS) ;
    return TRUE ; 
  }

  // Notify server that a process will be created
  // so it will not try to inject it in another manner.
  // The server response is 1 if we are allowed to hook.
  sdncp.hdr.dwCode	= SDN_CREATING_PROCESS ;
  sdncp.hdr.dwProcessId = GetCurrentProcessId () ;
  _tcslcpy (sdncp.szPath, tszPath, MAX_PATH) ;
  dwServerResponse = SpyInt_NotifyServer (&sdncp.hdr, sizeof(sdncp)) ;
  TRACE_INFO (TEXT("Server response : %u\n"), dwServerResponse) ;
  
  // provide a local PROCESS_INFORMATION if needed
  if( ! pProcessInformation )
    pProcessInformation = &pi ;

  // call real CreateProcessA
  bSuccess = pfnStub (szAppName, szCmdLine, pProcessAttr, pThreadAttr,
		      bInheritHandles, dwCreationFlags|CREATE_SUSPENDED, 
		      pEnvironment, szCurrentDirectory, pStartupInfo, 
		      pProcessInformation) ;
  
  // notify server if failed
  if( ! bSuccess )
    {
      dwSaveErrorCode = GetLastError () ;
      sdncp.hdr.dwCode = SDN_CREATING_FAILED ;
      SpyInt_NotifyServer (&sdncp.hdr, sizeof(sdncp)) ;
      free (tszCurDir) ;
      free (tszCmdLine) ;
      SetLastError (dwSaveErrorCode) ;
      return FALSE ;
    }

  // inject process if allowed
  if( dwServerResponse==1 )
    {
      if( ! SpyDll_InjectInSuspendedThread(pProcessInformation->hProcess, 
					   pProcessInformation->hThread) )
	TRACE_ERROR (TEXT("InjectInSuspendedThread failed\n")) ;     	  
    }

  // resume process if needed
  if( ! (dwCreationFlags&CREATE_SUSPENDED) ) 
    if( 1!=ResumeThread (pProcessInformation->hThread) )
      TRACE_ERROR (TEXT("Failed to resume thread")) ;
  
  // notify server
  sdnpc.hdr.dwCode	= SDN_PROCESS_CREATED ;
  sdnpc.hdr.dwProcessId	= GetCurrentProcessId () ;
  sdnpc.bHooked		= 0!=dwServerResponse ;
  sdnpc.dwProcessId	= pProcessInformation->dwProcessId ;
  _tcslcpy (sdnpc.szPath, tszPath, MAX_PATH) ;
  SpyInt_NotifyServer (&sdnpc.hdr, sizeof(sdnpc)) ;

  // now the unicode strings can be freed		  
  free (tszCurDir) ;
  free (tszCmdLine) ;

  return TRUE ;
}


/******************************************************************/
/* Exported function : CreateProcessW                             */
/******************************************************************/

BOOL WINAPI Spy_CreateProcessW (LPCWSTR			szAppName,
				LPWSTR			szCmdLine,
				LPSECURITY_ATTRIBUTES	pProcessAttr,
				LPSECURITY_ATTRIBUTES	pThreadAttr,
				BOOL			bInheritHandles,
				DWORD			dwCreationFlags,
				PVOID			pEnvironment,
				LPCWSTR			szCurrentDirectory,
				LPSTARTUPINFOW		pStartupInfo,
				LPPROCESS_INFORMATION	pProcessInformation) 
{
  PROC			pfnStub = Hooks_GetStubAddress (HOOKS_CREATEPROCESSW) ;

  DWORD			nReaction ;
  BOOL			bSuccess ;
  DWORD			dwServerResponse ;
  DWORD			dwSaveErrorCode ;
  SDNCREATINGPROCESS	sdncp ;
  SDNPROCESSCREATED	sdnpc ;
  PROCESS_INFORMATION	pi ;

  BOOL			bAppNameSpecified ;
  LPCTSTR		tszCmdLine ;
  TCHAR			tszPath[MAX_PATH] ;

  TRACE ;

  if( szAppName==NULL && szCmdLine==NULL )
    TRACE_WARNING (TEXT("szAppName==NULL and szCmdLine==NULL\n"));

  bAppNameSpecified = szAppName != NULL ;
  
  // get command line
  tszCmdLine = bAppNameSpecified ? szAppName : szCmdLine ;

  // extract program path 
  _HookSys_ExtractPathFromCmdLine (tszPath, szCurrentDirectory, tszCmdLine, 
				   !bAppNameSpecified, !bAppNameSpecified) ;

  // apply filter
  nReaction = SpyInt_CatchCall (REASON_SYS_EXECUTE, TEXT("s"), tszPath) ;

  // if filter reject call, just return false
  if( nReaction == RULE_REJECT ) {
    SetLastError (ERROR_ACCESS_DENIED) ;
    return FALSE ; 
  }

  // if filter reject call, just return true
  if( nReaction == RULE_FEIGN ) {
    SetLastError (ERROR_SUCCESS) ;
    return TRUE ;       
  }

  // Notify server that a process will be created
  // so it will not try to inject it in another manner.
  // The server response is 1 if we are allowed to hook.
  sdncp.hdr.dwCode	= SDN_CREATING_PROCESS ;
  sdncp.hdr.dwProcessId = GetCurrentProcessId () ;
  _tcslcpy (sdncp.szPath, tszPath, MAX_PATH) ;
  dwServerResponse = SpyInt_NotifyServer (&sdncp.hdr, sizeof(sdncp)) ;
  TRACE_INFO (TEXT("Server response : %u\n"), dwServerResponse) ;
 
  // provide a local PROCESS_INFORMATION if needed
  if( ! pProcessInformation )
    pProcessInformation = &pi ;

  // call real CreateProcessW
  bSuccess = pfnStub (szAppName, szCmdLine, pProcessAttr, pThreadAttr,
		      bInheritHandles, dwCreationFlags|CREATE_SUSPENDED, 
		      pEnvironment, szCurrentDirectory, pStartupInfo, 
		      pProcessInformation) ;
  
  // notify server if failed
  if( ! bSuccess )
    {
      dwSaveErrorCode = GetLastError () ;
      sdncp.hdr.dwCode = SDN_CREATING_FAILED ;
      SpyInt_NotifyServer (&sdncp.hdr, sizeof(sdncp)) ;
      SetLastError (dwSaveErrorCode) ;
      return FALSE ;
    }

  // inject process if allowed
  if( dwServerResponse==1 )
    {
      if( ! SpyDll_InjectInSuspendedThread(pProcessInformation->hProcess, 
					   pProcessInformation->hThread) )
	TRACE_ERROR (TEXT("InjectInSuspendedThread failed\n")) ;
    }
  
  // resume process if needed
  if( ! (dwCreationFlags&CREATE_SUSPENDED) ) 
    if( 1!=ResumeThread (pProcessInformation->hThread) )
      TRACE_ERROR (TEXT("Failed to resume thread")) ;
  
  // notify server
  sdnpc.hdr.dwCode	= SDN_PROCESS_CREATED ;
  sdnpc.hdr.dwProcessId	= GetCurrentProcessId () ;
  sdnpc.dwProcessId	= pProcessInformation->dwProcessId ;
  sdnpc.bHooked		= 0!=dwServerResponse ;
  _tcslcpy (sdnpc.szPath, tszPath, MAX_PATH) ;
  SpyInt_NotifyServer (&sdnpc.hdr, sizeof(sdnpc)) ;
   
  return bSuccess ;
}


/******************************************************************/
/* Internal function : ExtractPathFromCmdLine                     */
/******************************************************************/
  
BOOL _HookSys_ExtractPathFromCmdLine (LPTSTR szPath, LPCTSTR szCurDir, LPCTSTR szCmdLine, 
				      BOOL bRemoveArgs, BOOL bSearchPath)
{ 
  TCHAR szAppName[MAX_PATH] ;

  // verify params
  ASSERT (szPath!=NULL) ;
 
  //
  // Extract application from command line
  //
  if( bRemoveArgs )
    {
      LPCTSTR	pBegin, pEnd ;
      UINT	nLength ;
      
      if( szCmdLine[0]==TEXT('\"') ) 
	{
	  pBegin = szCmdLine + 1 ;
	  pEnd = _tcschr (pBegin, TEXT('\"')) ;
	}
      else
	{
	  pBegin = szCmdLine ;
	  pEnd = _tcschr (pBegin, TEXT(' ')) ;
	}

      ASSERT (pBegin!=NULL) ;
      
      if( ! pEnd )
	pEnd = pBegin + _tcslen(pBegin) ;
      
      nLength = pEnd-pBegin ;
      if( nLength>MAX_PATH-1 ) nLength = MAX_PATH-1 ;
      
      CopyMemory (szAppName, pBegin, nLength*sizeof(TCHAR)) ;
      szAppName[nLength] = 0 ;
    }
  else
    {
      PathCanonicalize (szAppName, szCmdLine) ;      
    }
 
  //
  // Is path specified ?
  //
  if( ! PathIsFileSpec (szAppName) )
    {
      _tcscpy (szPath, szAppName) ;
      
      if( PathFileExists(szPath) )
	return TRUE ;      

      // no extension => append ".exe"
      if( PathFindExtension(szPath)[0]==0 )
	{
	  PathAddExtension (szPath, NULL) ; 
      
	  if( PathFileExists(szPath) )
	    return TRUE ;  
	}

      return FALSE ;
    }
  
  //
  // 1. The directory from which the application loaded.
  //
  if( szCurDir )
    {
      PathCombine (szPath, szCurDir, szAppName) ;
      
      if( PathFileExists(szPath) )
	return TRUE ;  
      
      // no extension => append ".exe"
      if( PathFindExtension(szPath)[0]==0 )
	{
	  PathAddExtension (szPath, NULL) ;
	  
	  if( PathFileExists(szPath) )
	    return TRUE ;  
	}
    }   
  
  //
  // 2. The current directory for the parent process.
  // 
  
  GetCurrentDirectory (MAX_PATH, szPath) ;  
  PathCombine (szPath, szPath, szAppName) ;

  if( PathFileExists(szPath) )
    return TRUE ; 

  // no extension => append ".exe"		    
  if( PathFindExtension(szPath)[0]==0 )
    {
      PathAddExtension (szPath, NULL) ;
      if( PathFileExists(szPath) )
	return TRUE ;
    }

  //
  // 3. The 32-bit Windows system directory.
  //
  GetSystemDirectory (szPath, MAX_PATH) ;
  PathCombine (szPath, szPath, szAppName) ;

  if( PathFileExists(szPath) )
    return TRUE ; 

  // no extension => append ".exe"		    
  if( PathFindExtension(szPath)[0]==0 )
    {
      PathAddExtension (szPath, NULL) ;
      if( PathFileExists(szPath) )
	return TRUE ;
    }

  //
  // 4. The 16-bit Windows system directory. (ignored in this implementation)
  //

  //
  // 5. The Windows directory.
  //
  GetWindowsDirectory (szPath, MAX_PATH) ;
  PathCombine (szPath, szPath, szAppName) ;

  if( PathFileExists(szPath) )
    return TRUE ; 

  // no extension => append ".exe"		    
  if( PathFindExtension(szPath)[0]==0 )
    {
      PathAddExtension (szPath, NULL) ;
      if( PathFileExists(szPath) )
	return TRUE ;
    }  
  
  //
  // 6. The directories that are listed in the PATH environment variable.
  //

  if( bSearchPath )
    {
      LPTSTR	szPathEnv ;
      UINT	nPathEnvLen ;
      LPTSTR	szCurDir ;
      int	i ;

      nPathEnvLen = GetEnvironmentVariable (TEXT("PATH"), NULL, 0) ;
      szPathEnv = malloc (nPathEnvLen*sizeof(TCHAR)) ;
      
      GetEnvironmentVariable (TEXT("PATH"), szPathEnv, nPathEnvLen) ;

      szCurDir = szPathEnv ;
      for( i=0 ; i<32767 ; i++ )
	{
	  if( szPathEnv[i]==0 )
	    {
	      if( szCurDir[0] )
		{
		  PathCombine (szPath, szCurDir, szAppName) ;

		  if( PathFileExists(szPath) )
		    return TRUE ; 
		  
		  // no extension => append ".exe"		    
		  if( PathFindExtension(szPath)[0]==0 )
		    {
		      PathAddExtension (szPath, NULL) ;
		      if( PathFileExists(szPath) )
			return TRUE ;
		    }  
		}
	      break ;
	    }

	  if( szPathEnv[i]==TEXT(';') )
	    {
	      szPathEnv[i] = 0 ;

	      if( szCurDir[0] )
		{
		  PathCombine (szPath, szCurDir, szAppName) ;
		  
		  if( PathFileExists(szPath) )
		    return TRUE ; 
		  
		  // no extension => append ".exe"		    
		  if( PathFindExtension(szPath)[0]==0 )
		    {
		      PathAddExtension (szPath, NULL) ;
		      if( PathFileExists(szPath) )
			return TRUE ;
		    }  
		}

	      szCurDir = &szPathEnv[i+1] ;
	    }
	}	
    }

  _tcscpy (szPath, szAppName) ;

  return FALSE ;  
}
